#pragma once

#include <cmath>

#include "base.h"
#include "../impl/ma.h"
#include "../impl/s2.h"

namespace mdtx::online::indicator
{
    template <typename ScImpl>
    class zscore_wrapper : public indi_base<typename ScImpl::scala_in,
                                            typename ScImpl::scala_out,
                                            ScImpl::num_in, 1>
    {
        using base = indi_base<typename ScImpl::scala_in,
                               typename ScImpl::scala_out,
                               ScImpl::num_in, 1>;
        using scale_indi_type = std::enable_if_t<ScImpl::num_out == 2, ScImpl>;

    public:
        using typename base::scala_out;
        using typename base::tuple_in;
        using typename base::tuple_out;
        using base::update;

    protected:
        scale_indi_type scale;

    public:
        template <typename... Args>
        explicit zscore_wrapper(Args... args)
            : scale(args...)
        {
        }

        tuple_out update(tuple_in const &in) override
        {
            scala_out val{};
            auto [m, s] = scale.update(in);
            if (s > 0)
            {
                val = (std::get<0>(in) - m) / s;
            }
            return tuple_out(val);
        }
    };

    template <typename ScImpl>
    class smoothz_wrapper : public indi_base<typename ScImpl::scala_in,
                                             typename ScImpl::scala_out,
                                             ScImpl::num_in, 1>
    {
        using base = indi_base<typename ScImpl::scala_in,
                               typename ScImpl::scala_out,
                               ScImpl::num_in, 1>;
        using scale_indi_type = std::enable_if_t<ScImpl::num_out == 2, ScImpl>;

    public:
        using typename base::scala_in;
        using typename base::scala_out;
        using typename base::tuple_in;
        using typename base::tuple_out;
        using base::update;

    protected:
        scale_indi_type scale;
        ptrdiff_t lag, n{};
        scala_in z, r, last{};
        scala_out m{}, s{};

    public:
        template <typename... Args>
        explicit smoothz_wrapper(ptrdiff_t lag, scala_in z, scala_in r, Args... args)
            : scale(args...),
              lag(lag), z(z), r(r)
        {
        }

        tuple_out update(tuple_in const &in) override
        {
            tuple_in in_copy{in};
            scala_in x = std::get<0>(in);
            scala_out signal{};

            if (n == lag)
            {
                if (s > 0 && std::abs(x - m) > z * s)
                {
                    signal = (x - m) / s;
                    x = r * x + (1 - r) * last;
                }
            }
            else
            {
                n += 1;
            }
            std::get<0>(in_copy) = last = x;
            std::tie(m, s) = scale.update(in_copy);

            return tuple_out(signal);
        }
    };

    template <typename T>
    class smoothz : indi_unary1<T>
    {
        using base = indi_unary1<T>;

    public:
        using typename base::tuple_in;
        using typename base::tuple_out;
        using base::update;

    protected:
        boost::circular_buffer<T> buf, bufm;
        impl::sma<T> ma{};
        impl::var<T> var{};
        T z, r;

    public:
        smoothz(ptrdiff_t period, ptrdiff_t m_period, T z, T r)
            : buf(period), bufm(m_period),
              z(z), r(r)
        {
        }

        tuple_out update(tuple_in const &in) override
        {
            T x = std::get<0>(in);
            T x_atten{x}, signal{};

            // s2
            if (buf.full())
            {
                T sd = std::sqrt(var.s2 / (var.n - 1));
                if (sd > 0 && std::abs(x - var.m) > z * sd)
                {
                    signal = (x - var.m) / sd;
                    x_atten = r * x + (1 - r) * ma.m;
                }
                var.remove_impl(buf.front());
            }
            var.insert_impl(x_atten);
            buf.push_back(x_atten);

            // ma
            if (bufm.full())
            {
                ma.roll_impl(x, bufm.front());
            }
            else
            {
                ma.insert_impl(x);
            }
            bufm.push_back(x);

            return tuple_out(signal);
        }
    };

    template <typename T>
    class smoothz_cu : indi_unary1<T>
    {
        using base = indi_unary1<T>;

    public:
        using typename base::tuple_in;
        using typename base::tuple_out;
        using base::update;

    protected:
        ptrdiff_t lag, n{};
        boost::circular_buffer<T> bufm;
        impl::sma<T> ma{};
        impl::var<T> var{};
        T z, r;

    public:
        smoothz_cu(ptrdiff_t lag, ptrdiff_t m_period, T z, T r)
            : lag(lag),
              bufm(m_period),
              z(z), r(r)
        {
        }

        tuple_out update(tuple_in const &in) override
        {
            T x = std::get<0>(in);
            T x_atten{x}, signal{};

            // s2
            if (n == lag)
            {
                T sd = std::sqrt(var.s2 / (var.n - 1));
                if (sd > 0 && std::abs(x - var.m) > z * sd)
                {
                    signal = (x - var.m) / sd;
                    x_atten = r * x + (1 - r) * ma.m;
                }
            }
            else
            {
                n += 1;
            }
            var.insert_impl(x_atten);

            // ma
            if (bufm.full())
            {
                ma.roll_impl(x, bufm.front());
            }
            else
            {
                ma.insert_impl(x);
            }
            bufm.push_back(x);

            return tuple_out(signal);
        }
    };

    template <typename T>
    class smoothz_ew : indi_unary1<T>
    {
        using base = indi_unary1<T>;

    public:
        using typename base::tuple_in;
        using typename base::tuple_out;
        using base::update;

    protected:
        ptrdiff_t lag, n{};
        impl::var_ew<T> var;
        impl::ema<T> ma;
        T z, r;

    public:
        smoothz_ew(ptrdiff_t period, ptrdiff_t m_period, T z, T r)
            : lag(period),
              var(T(2) / (period + 1)),
              ma(T(2) / (m_period + 1)),
              z(z), r(r)
        {
        }

        tuple_out update(tuple_in const &in) override
        {
            T x = std::get<0>(in);
            T x_atten{x}, signal{};

            // s2
            if (n == lag)
            {
                T sd = std::sqrt(var.s2);
                if (sd > 0 && std::abs(x - var.m) > z * sd)
                {
                    signal = (x - var.m) / sd;
                    x_atten = r * x + (1 - r) * ma.m;
                }
            }
            else
            {
                n += 1;
            }
            var.insert_impl(x_atten);

            // ma
            ma.insert_impl(x);

            return tuple_out(signal);
        }
    };
}
